#!/usr/bin/env python
import pypjl,pfsScanner
import random, pickle, sys, optparse, os, re, hashlib, bz2
try:
    from Crypto.Cipher import AES
except ImportError:
    print "You need to install pycrypto"
    exit()
#testing
import pdb
class crypto:
    def encrypt(self,file):
        """
        Encrypt a file with a random 256bit key and iv
        file = file to encrypt
        function returns encrypted file, random iv, and random key for when
        it is needed to decrypt later
        """
        print "Encrypting File"
        random.seed(os.urandom(random.randrange(5000,10000)))
        ranIV = ''.join(FileSys().genRanStr(16))
        #cycle the pseudo RNG a buch before we generate the key
        z = FileSys().genRanStr(random.randrange(1000,2000))
        for y in z:
            FileSys().genRanStr(random.randrange(101,153))
        ranKey = ''.join(FileSys().genRanStr(32))
        encObj = AES.new(ranKey, AES.MODE_CBC, ranIV)
        listobj = [file]
        baseLen = len(file)
        nounce = ''.join(listobj)
        modNounce = len(nounce)%16
        if modNounce == 0:
            ciph = encObj.encrypt(nounce)
        else: 
            ciph = encObj.encrypt(nounce+''.join(FileSys().genRanStr(16-modNounce)))
        
        print "File was Encrypted"
        return [{"iv":ranIV,"key":ranKey,"ustrLen":baseLen},ciph]
    
    def decrypt(self,iv,key,encstr,baselen):
        """
        decrypt a given string
        iv = 128bit iv used to encrypt
        key = 256bit key used to encrypt
        encstr = encrypted string to decrypt
        baselen = orginal len of string before padding
        returns an un encrypted string
        """
        encObj = AES.new(key, AES.MODE_CBC, iv)
        unciph = encObj.decrypt(encstr)
        return unciph[:baselen]

class FileSys:
    def __init__(self):
        #enable random comments
        self.plib     = pypjl.PJLDynamic(True)
        self.port     = 9100 
        self.plocateS = {} # serialnumber to IP lookup
        self.plocateI = {} # ip to serialnumber lookup 
        self.vpfi     = {} # virtual printer information table
        self.vpft     = {} # virtual printer File table
        self.data     = [] #used for dbug to dump all raw file info
        self.mdata    = [] #used for dbug to dump all raw model info
        self.randata  = None #random data used for uploading test files

    def exitFS(self,file="printFS.dat"):
        """
        write out the vpft and vpfi for use later
        file = the dat file where vpft and vpfi are saved to disk
        when program exits...
        """
        if self.vpfi != {}:
            try:os.remove(file) #delete old vpfi file
            except OSError: pass
            db = open(file,'w')
            pickle.dump([a.vpfi,a.vpft,a.plocateS,a.plocateI],db) #write new one
            db.close()
            print "\nvpft data written to",file

    def loadFS(self,file="printFS.dat"):
        """
        load the vpft for file acess
        file = the dat file where vpft and vpfi are saved to disk
        """
        db = open(file,'r')
        list = pickle.load(db)
        self.vpfi = list[0] #printer info
        self.vpft = list[1] #file info
        self.plocateS = list[2] #serial to IP
        self.plocateI = list[3] #IP to serial
        db.close()
        print "Loaded vpft data from",file

    def delDevice(self,ip):
        """
        remove a device from the vpft
        ip = device ip that you wish to remove to the vpfi
        """
        #should prob also go and remove all links to files that may be stored on this device
        serial = self.plocateI[ip]
        del self.vpfi[serial]
        print ip,"removed from the VPFT"

    def panicMode(self):
        """
        Reboot all printers and remove all evidence
        """
        x = raw_input("This will remove all data on all devices are you sure you want to do this? N\y ")
        if x.lower() == "y":
            y = raw_input("Are you really sure you want to remove ALL data from ALL devices? N\y ")
            if y.lower() == "y":
                print "DEFCON BLACK!\n"*5
                for ip in self.vpfi.values():
                    print "Rebooting",ip['ip']
                    #print self.plib.reboot(ip['ip'],self.port) #used for debug
                self.plocateS = {} # serialnumber to IP lookup
                self.plocateI = {} # ip to serialnumber lookup
                self.vpfi     = {} # virtual printer information table
                self.vpft     = {} # virtual printer File table
                try:os.remove('printFS.dat') #delete old vpfi file
                except OSError: pass
            else: print "Self Destruct Adverted!"
        else: print "Self Destruct Adverted!"
        sys.exit(0)

    def processIP(self,ipList):
        """
        Call the threaded scanner then parse out its returned data
        Populates self.vpfi and provides feedback from scanner
        """
        scanner = pfsScanner.scanEngine(ipList)
        #get scan data
        data = scanner.main()
        pdata = data[0]
        timer = data[1]
        rescan = data[2]
        if len(rescan) != 0:
            print "Rescanning %s devices due to download or hash check failures" % len(rescan)
            scanner2 = pfsScanner.scanEngine(rescan)
            data2 = scanner2.main()
            pdata2 = data[0]
            timer2 = data[1]
            timer += timer2
            if len(pdata2) != 0:
                pdata.update(pdata2)
                
        #loop through and populate vpfi
        for ip in pdata:
            #raw data
            rdata = pdata[ip]
            if rdata['supported'] is not False:
                self.plocateS = {rdata['scanSerial']:ip} # serialnumber to IP lookup
                self.plocateI = {ip:rdata['scanSerial']} # ip to serialnumber lookup
                self.vpfi[rdata['scanSerial']] = {
                    'ip':ip,
                    'serial':rdata['scanSerial'],
                    'label':rdata['scanDiskInfo']['label'],
                    'disknum':rdata['scanDiskInfo']['volume'],
                    'maxspace':rdata['scanDiskInfo']['totalSize'],
                    'freespace':rdata['scanDiskInfo']['freeSpace'],
                    'port':rdata['port'],
                    'model':rdata['scanModel'],
                    'ft':[]} #empty list of files stored here
                print "\nIP",ip,"is supported"
            else: 
                print "\nIP",ip,"is not supported"
        print "\n",len(self.vpfi.keys()),"Devices are Supported"
        print "\nScan Completed in",timer / 60,"Minutes\n"
        self.fstats('-m')

    def buildList(self,filename):
        """
        take a list of ips and build list of printers
        that it can use for a file system
        filename = file to read ip's from, ips should be one per line
        """
        f = open(filename,"r")
        lines  = f.readlines()
        #remove commented lines
        for ip in lines:
            if ip[0] == "#":
                lines.remove(ip)
                continue
            else: continue
        print "Building Table for",len(lines),"Devices\nThis may take a bit\n"
        self.processIP(lines)
        #save the VPFT to disk
        a.exitFS()
        return
    
    def genRanStr(self,bytelen):
        """
        generate a random string of data
        byelen = number of bytes to generate
        """
        letters = [] #empty list for mixed letters
        random.seed() #use system time as seed
        chlist = []
        while len(letters) < bytelen:
            #generate some sudo random capital letters
            chlist.append(random.randrange(66,90,1))
            chlist.append(random.randrange(97,122,1))
            letters.append(chr(random.choice(chlist)))
        #join both lists and shuffle the letters
        random.shuffle(letters)
        return letters

    def genFname(self):
        """
        generate a random filename between 1 and 40 bytes
        should mean 4.36650282 x 10^68 possible unique file names
        """
        #return a random string name with random length
        return ''.join(random.sample(self.genRanStr(40),random.randint(1,40)))

    def getStrIp(self,fileSize):
        """
        figure out the best two printers to store a given file on
        returns a list of two ip addresses
        """
        #ideas
        #some sort of cost calculation that gives the best devices 
        #based on their amount of free space the safty of the data storage 
        #and the size of the file that can be upload to them
        #this is ineffienct and need to find a way 
        #to do it with out sorting though vpfi twice
        ip = {} #ips to store files on
        sizes=[]
        for printer in self.vpfi.values():
            sizes.append(int(printer['freespace']))
        if len(sizes) > 1:
            sizes.sort() #sort so smallest num first
        x = []
        for size in sizes:
            if fileSize < size:
                x.append(size)
            if len(x) == 2: break
        for printer in self.vpfi.values():
            if int(printer['freespace']) in x:
                ip[printer['serial']] = printer['ip']
        return ip 
 
    def fstats(self,form="b"):
        """
        calculate and print file stats
        """
        mspace = 0 #maxspace
        fspace = 0 #freespace
        for printer in self.vpfi:
            mspace += float(self.vpfi[printer]["maxspace"])
            fspace += float(self.vpfi[printer]["freespace"])
        if form.lower() not in ["-b","-k","-m","-g"]:
            print "Invlaid Space Format request",form
            print "Using Bytes as default"
            form = "-b"
        if form.lower() == "-b":
            print "Total Space",int(mspace),"Bytes"
            print "Free Space",int(fspace),"Bytes\n"
        elif form.lower() == "-k":
            print "Total Space",int(mspace)/1024,"Kilobytes"
            print "Free Space",int(fspace)/1024,"KiloBytes\n"
        elif form.lower() == "-m":
            print "Total Space",(int(mspace)/1024)/1024,"Megabytes"
            print "Free Space",(int(fspace)/1024)/1024,"MegaBytes\n"
        elif form.lower() == "-g":
            print "Total Space",((int(mspace)/1024)/1024)/1024,"Gigabytes"
            print "Free Space",((int(fspace)/1024)/1024)/1024,"GigaBytes\n"

        try:x = fspace / mspace * 100 
        except ZeroDivisionError: x= 0
        print "Percent Free Space",str(int(x))+"%\n"
        print len(self.vpfi.keys()),"Printers Support File Storage" 

class fileCommands:
    """
    Class for comands for dealing with the virtual filesystem
    """
    def __init__(self):
        self.plib   = pypjl.PJLDynamic(True)
        self.pwd    = "\\" #should really get what ever root 
        #is for the given device we are working with
    
    def dirlist(self):
        """
        print out a list of the file in a given virtual directory
        dir = directory path you with to print #future addition
        """
        for file in a.vpft.keys():
            print file
    
    def chdir(self,gdir):
        """
        not used at this time....
        change directory
        gdir = directory path you wish to change to
        """
        self.pwd = gdir

    def fdcopy(self,file):
        """
        copy a file down from the printer to the host
        """
        if file in a.vpft.keys():
            for device in a.vpft[file]:
                filepath = a.vpft[file][device]['filepath']
                serial   = device #serial number
                sha      = a.vpft[file][device]['sha']
                ip       = a.vpfi[serial]['ip']
                sfName   = a.vpft[file][device]['sfName']
                comp     = a.vpft[file][device]['bz2']
                crypt    = a.vpft[file][device]['crypto']
                try:
                    x = self.plib.downloadFiles(filepath+sfName,ip,a.port)
                    if x != False:
                        #decrypt
                        print "Decrypting File"
                        x = crypto().decrypt(crypt['iv'],crypt['key'],x,crypt['ustrLen'])
                        #uncompress data stream
                        if comp == True: z = bz2.decompress(x)
                        elif comp == False: z = x
                        if hashlib.sha256(z).hexdigest() != sha:
                            print "Error with file "+name
                            print "Sha224 Error! hashes do not match"
                            print "Check your received file"
                            print "Still writing file"
                        self.fileWrite(z,file)
                        print "Retrived File",file
                        return
                    else: continue
                except Exception:
                    #continue #if the download files it should try the next storage device then update the vpfi and vpft
                    pass
        else: print file,"not found!"
        print "Unable to retrive file",file
        return 
    
    def fileWrite(self,data,filename):
        """
        Open a file and write out data recived from a printer
        filename = filename and path to write to
        """
        try:
            f = open(filename,'w')
            f.write(data)
            f.close()
        except IOError,e:
            print e
            return False
    
    def fileRead(self,file):
        """
        Open a file read it and return it as a string
        file = filename and path to read from
        """
        try:
            f = open(file,'r')
            x = f.read()
            f.close()
            return x
        except IOError,e:
            print e
            return False

    def fucopy(self,file,filepath=""):
        """
        copy a file up from the host to the printer
        file filepath of file you wish to upload
        """
        data = self.fileRead(file)
        if data == False:
            return
        #get intial sha before upload and while uncompressed
        sha = hashlib.sha256(data).hexdigest()
        cdata = bz2.compress(data)#compress data    
        fsize = len(cdata)
        #check to see if we can compres the file if not dont compress
        if fsize >= len(data):
            fsize = len(data)
            state = False
            udata = data #upload data
        else:
            udata = cdata #upload data
            state = True
            try:dd = float(len(cdata))/float(len(data))*100
            except ZeroDivisionError: x= 0
            print "File Compressed by",str(dd)+"%\n"
        #add encryption here
        cryptRtn = crypto().encrypt(udata)
        udata = cryptRtn[1]
        hosts = a.getStrIp(fsize)
        for serial in hosts.keys():
            if filepath == "": filepath = a.vpfi[serial]["disknum"]+"\\"
            ip = hosts[serial] #serial key 
            sfName = a.genFname()
            x = self.plib.uploadFiles(udata,filepath+sfName,ip,a.port) #a.port = 9100
            if x == False: 
                print file,"Failed to upload"
                return
            #in all reality this freespace update should be a 
            #function and return an error if there is not enough space!
            a.vpfi[serial]['freespace'] = str(
                int(a.vpfi[serial]['freespace']) - int(fsize)
                )
            #check if file allready exists then update 
            #postion where both items are stored. use a nested list
            if file not in a.vpfi[serial]['ft']:
                #append the file name to the list of files on this device
                a.vpfi[serial]['ft'].append(file) 
            else: 
                print "Overwriting orginal copy of",file #should check to see if file allready exists...
            if file in a.vpft:
                    a.vpft[file][serial] = {
                        "filepath":filepath,"filename":file,
                        "sha":sha,"size":fsize,
                        "sfName":sfName, #random file name used to store file
                        "bz2":state,
                        "crypto":cryptRtn[0]}
            else:
                a.vpft[file] = {serial:
                    {"filepath":filepath,
                    "filename":file,
                    "sha":sha,
                    "size":fsize,
                    "sfName":sfName,
                    "bz2":state,
                    "crypto":cryptRtn[0]}}

            #print a.vpfi[serial]['ft'] #for debug
            #print a.vpft #for debug
        print "Stored File",file
        #need some way of querying the space knowing 
        #what printer has it and choosing the file to upload to
    
    def fdel(self,file):
        """
        delete a file
        """
        if file in a.vpft.keys():
            for device in a.vpft[file].keys(): #device is serial#
                filepath = a.vpft[file][device]['filepath']
                sfName = a.vpft[file][device]['sfName']
                serial = device
                ip = a.vpfi[serial]['ip']
                self.plib.delete(filepath+sfName,ip,a.port) #a.port = 9100
                a.vpfi[serial]['ft'].remove(file) #remove from vpfi db
                a.vpfi[serial]['freespace'] = str(
                    long(a.vpfi[serial]['freespace']) + long(
                        a.vpft[file][device]["size"])) #update freespace
            del a.vpft[file] #remove file from vpft db
            print "Removed",file,"from the file table."

    def addip(self,ip):
        """
        add a single printer to the vpft
        """
        ip = ip.strip()
        a.processIP([ip])
    
    def dumpDevices(self):
        """
        Print out supported deviecs and their model numbers
        """
        for x in a.vpfi.values(): 
            print '%s,%s' %(x['model'],x['ip'])

class userIface:
    """
    class for taking input from the user on the cmd line
    """
    def __init__(self):
        self.g = fileCommands()
        self.cmdList = {"put":[self.g.fucopy,"Usage: put FILENAME"],
            "get":[self.g.fdcopy,",Usage: get FILENAME"],
            "ls":[self.g.dirlist,",Usage: ls "],
            "addip":[self.g.addip,",Usage: addip IPADRESS"],
            "help":[self.cmdHelp,",Usage: help"],
            "rm":[self.g.fdel,",Usage: rm FILENAME"],
            "stats":[a.fstats,",Usage: stats [-b -k -m -g]"],
            "panic":[a.panicMode,",Usage: panic"],
            "rmip":[a.delDevice,",Usage: rmip IPADRESS"],
            "write":[a.exitFS,",Usage: write out the vpft to the disk"], #write out the vpft without needing to exit
            "devices":[self.g.dumpDevices,",Usage: return supported devices and their IP addresses"],
            "pdb":[pdb.set_trace,",Usage: pdb"]}

    def cmdHelp(self):
        """
        loop though and print interactive commands
        """
        print "Commands are ..."
        for item in self.cmdList.keys():
            print item, self.cmdList[item][1]

    def getInput(self):
        """
        return user input in lowercase and split into a list
        """
        return raw_input("\nFile Command? >>> ").split(' ')

    def callcmd(self):
        """
        parse user commands and determine what to do
        """
        usrcmd = self.getInput()
        usrcmd[0] = usrcmd[0].lower() #make cmd lowercase
        if usrcmd[0] == "exit":
            a.exitFS(options.data)
            sys.exit(0)
        if usrcmd[0] in self.cmdList.keys():
            try:
                try:
                    return self.cmdList[usrcmd[0]][0](usrcmd[1])
                except IndexError:
                    return self.cmdList[usrcmd[0]][0]()
                except TypeError:
                    return self.cmdList[usrcmd[0]][0](usrcmd[1],usrcmd[2])
            except AttributeError:
                return self.cmdList[usrcmd[0]][0]()
            except IndexError:
                print "\nInvalid Command, Use the help command to see all commands"
            except TypeError:
                print "\nInvalid Command, Use the help command to see all commands"
        else:
            print "\nInvalid Command, Use the help command to see all commands"
            return

    def iface(self):
        """
        take userinput on a loop and run it as a command
        """
        while True:
            self.callcmd()

if __name__ == "__main__":
    parser = optparse.OptionParser("usage: %prog options -h -i -l")
    parser.add_option("-i", "--interactive", dest="int", action="store_true",
    default=False, help="Use interactive mode")
    parser.add_option("-d", "--dataFile", dest="data", nargs=1,
    default="printFS.dat", help="Load up a differnt file table data file")
    parser.add_option("-l", "--list", dest="list",nargs=1, 
        default=None ,help="List of Ips to use")
    if len(sys.argv) <= 1:
        parser.print_help()
        sys.exit(0)
    (options, args) = parser.parse_args()
    if options.list in ['-i','--interactive'] :
        print "Error you can not use -l and -i or the same time"
        parser.print_help()
        sys.exit(-1)
    a = FileSys()
    b = userIface()
    #check to see if a vpft exits
    if options.list != None:
        if os.path.exists(options.data):
            print "VPFT Exists\n"
            x = None
            while x not in ['n','y']:
                x = raw_input("Do you want to overwrite your current file table?\nYOU WILL LOSE ALL STORED FILES! n\y?").lower()
                if x == 'y':
                    os.remove(options.data)
                    a.buildList(options.list)
                    b.iface()
                elif x == 'n':
                    a.loadFS(options.data)
                    b.iface()
        else:
            a.buildList(options.list)
            b.iface()
    
    elif options.int == True: #prob not needed
        if os.path.exists(options.data) == True:
            a.loadFS(options.data)
            b.iface()
        else:
            print "Error VPFT %s Not Found" % options.data
            b.iface()
